from nlp_tools.utils.model import convert_to_saved_model
from nlp_tools.tasks.labeling import ABCLabelingModel
import os
import requests
import numpy as np
from nlp_tools.server_utils.tfserver_utils.base_tfserver_utils import BaseTfServerUtils

model_save_path = '/home/qiufengfeng/nlp/train_models/ner'
tfserver_model= BaseTfServerUtils(model_save_path)

text = "孙高涵和他的朋友吴柏龙在一起玩"

tensor = tfserver_model.get_tfserver_inputs([text])

def extract_labels(text,ners):
    ner_reg_list = []
    if ners :
        new_ners = []
        for ner in ners:
            new_ners += ner

        for word,tag in zip([char for char in text],new_ners):
            if tag != "O":
                ner_reg_list.append((word,tag))

    # 输出模型的NER识别结果
    labels = {}
    if ner_reg_list:
        for i,item in enumerate(ner_reg_list):
            if item[1].startswith('B'):
                label = ""
                end = i + 1
                while end <= len(ner_reg_list) - 1 and ner_reg_list[end][1].startswith('I'):
                    end += 1

                ner_type = item[1].split("-")[1]

                if ner_type not in labels.keys():
                    labels[ner_type] = []

                label += "".join([item[0] for item in ner_reg_list[i:end]])
                labels[ner_type].append(label)
    return labels


# predict
r = requests.post("http://localhost:8501/v1/models/company_name_entity:predict", json={"instances": tensor})
preds = r.json()['predictions']
predict = np.array(preds).argmax(-1)
labels = label_processor.inverse_transform(predict)
labels = extract_labels(segment_text, labels)
print(labels)

'''
curl -d '{"Input-Segment": [1.0, 2.0, 5.0]}' \
   -X POST http://localhost:8501/v1/models/classification:predict'''
